// Copyright 2019 JD.com Inc. JD AI

// Tencent is pleased to support the open source community by making ncnn
// available.
//
// Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved.
//
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

#include "allocator.h"

#include <stdio.h>

namespace ncnn {

PoolAllocator::PoolAllocator() {
    size_compare_ratio = 192;  // 0.75f * 256
}

PoolAllocator::~PoolAllocator() {
    clear();

    if (!payouts.empty()) {
        fprintf(stderr, "FATAL ERROR! pool allocator destroyed too early\n");
        std::list<std::pair<size_t, void *>>::iterator it = payouts.begin();
        for (; it != payouts.end(); it++) {
            void *ptr = it->second;
            fprintf(stderr, "%p still in use\n", ptr);
        }
    }
}

void PoolAllocator::clear() {
    budgets_lock.lock();

    std::list<std::pair<size_t, void *>>::iterator it = budgets.begin();
    for (; it != budgets.end(); it++) {
        void *ptr = it->second;
        ncnn::fastFree(ptr);
    }
    budgets.clear();

    budgets_lock.unlock();
}

void PoolAllocator::set_size_compare_ratio(float scr) {
    if (scr < 0.f || scr > 1.f) {
        fprintf(stderr, "invalid size compare ratio %f\n", scr);
        return;
    }

    size_compare_ratio = (unsigned int)(scr * 256);
}

void *PoolAllocator::fastMalloc(size_t size) {
    budgets_lock.lock();

    // find free budget
    std::list<std::pair<size_t, void *>>::iterator it = budgets.begin();
    for (; it != budgets.end(); it++) {
        size_t bs = it->first;

        // size_compare_ratio ~ 100%
        if (bs >= size && ((bs * size_compare_ratio) >> 8) <= size) {
            void *ptr = it->second;

            budgets.erase(it);

            budgets_lock.unlock();

            payouts_lock.lock();

            payouts.push_back(std::make_pair(bs, ptr));

            payouts_lock.unlock();

            return ptr;
        }
    }

    budgets_lock.unlock();

    // new
    void *ptr = ncnn::fastMalloc(size);

    payouts_lock.lock();

    payouts.push_back(std::make_pair(size, ptr));

    payouts_lock.unlock();

    return ptr;
}

void PoolAllocator::fastFree(void *ptr) {
    payouts_lock.lock();

    // return to budgets
    std::list<std::pair<size_t, void *>>::iterator it = payouts.begin();
    for (; it != payouts.end(); it++) {
        if (it->second == ptr) {
            size_t size = it->first;

            payouts.erase(it);

            payouts_lock.unlock();

            budgets_lock.lock();

            budgets.push_back(std::make_pair(size, ptr));

            budgets_lock.unlock();

            return;
        }
    }

    payouts_lock.unlock();

    fprintf(stderr, "FATAL ERROR! pool allocator get wild %p\n", ptr);
    ncnn::fastFree(ptr);
}

UnlockedPoolAllocator::UnlockedPoolAllocator() {
    size_compare_ratio = 192;  // 0.75f * 256
}

UnlockedPoolAllocator::~UnlockedPoolAllocator() {
    clear();

    if (!payouts.empty()) {
        fprintf(stderr,
                "FATAL ERROR! unlocked pool allocator destroyed too early\n");
        std::list<std::pair<size_t, void *>>::iterator it = payouts.begin();
        for (; it != payouts.end(); it++) {
            void *ptr = it->second;
            fprintf(stderr, "%p still in use\n", ptr);
        }
    }
}

void UnlockedPoolAllocator::clear() {
    std::list<std::pair<size_t, void *>>::iterator it = budgets.begin();
    for (; it != budgets.end(); it++) {
        void *ptr = it->second;
        ncnn::fastFree(ptr);
    }
    budgets.clear();
}

void UnlockedPoolAllocator::set_size_compare_ratio(float scr) {
    if (scr < 0.f || scr > 1.f) {
        fprintf(stderr, "invalid size compare ratio %f\n", scr);
        return;
    }

    size_compare_ratio = (unsigned int)(scr * 256);
}

void *UnlockedPoolAllocator::fastMalloc(size_t size) {
    // find free budget
    std::list<std::pair<size_t, void *>>::iterator it = budgets.begin();
    for (; it != budgets.end(); it++) {
        size_t bs = it->first;

        // size_compare_ratio ~ 100%
        if (bs >= size && ((bs * size_compare_ratio) >> 8) <= size) {
            void *ptr = it->second;

            budgets.erase(it);

            payouts.push_back(std::make_pair(bs, ptr));

            return ptr;
        }
    }

    // new
    void *ptr = ncnn::fastMalloc(size);

    payouts.push_back(std::make_pair(size, ptr));

    return ptr;
}

void UnlockedPoolAllocator::fastFree(void *ptr) {
    // return to budgets
    std::list<std::pair<size_t, void *>>::iterator it = payouts.begin();
    for (; it != payouts.end(); it++) {
        if (it->second == ptr) {
            size_t size = it->first;

            payouts.erase(it);

            budgets.push_back(std::make_pair(size, ptr));

            return;
        }
    }

    fprintf(stderr, "FATAL ERROR! unlocked pool allocator get wild %p\n", ptr);
    ncnn::fastFree(ptr);
}

}  // namespace ncnn
